//
// Created by kai-xu on 2023/2/28.
//

#include "preprocess/matrixvt.h"

namespace autodrive {
namespace perception {
int32_t cal_coor_cpu(CustomCoorPoint *key_value, double *aug_mat,
                     double *ego_mat, std::vector<float> const &voxel_size,
                     std::vector<float> const &coors_range, int32_t sd_x,
                     int32_t sd_y, int32_t n_cam, int32_t dz, int32_t width,
                     int32_t height) {
  const float voxel_x = voxel_size[0];
  const float voxel_y = voxel_size[1];
  const float voxel_z = voxel_size[2];
  const float coors_x_min = coors_range[0];
  const float coors_y_min = coors_range[1];
  const float coors_z_min = coors_range[2];

  auto par_x =
      static_cast<float>((sd_x * width - 1.0) / (width - 1.0));  // 16.3488
  auto par_y =
      static_cast<float>((sd_y * height - 1.0) / (height - 1.0));  // 17

  for (int32_t k{0}; k < n_cam * dz; ++k) {
    for (int32_t i{0}; i < height; ++i) {
      for (int32_t j{0}; j < width; ++j) {
        auto mat_id = int32_t(k / dz);
        auto dz_id = float(k % dz * 0.45);
        float dz_index = dz_id + 1;
        float i_ = float(i) * par_y;
        float j_ = float(j) * par_x;

        double *aug_mat_add = aug_mat + mat_id * 16;
        double d2 = aug_mat_add[8] * j_ + aug_mat_add[9] * i_ +
                    aug_mat_add[10] * dz_index + aug_mat_add[11];
        double w2 = (aug_mat_add[0] * j_ + aug_mat_add[1] * i_ +
                     aug_mat_add[2] * dz_index + aug_mat_add[3] * 1) *
                    d2;
        double h2 = (aug_mat_add[4] * j_ + aug_mat_add[5] * i_ +
                     aug_mat_add[6] * dz_index + aug_mat_add[7] * 1) *
                    d2;

        double *ego_mat_add = ego_mat + mat_id * 16;
        double d3 = ego_mat_add[8] * w2 + ego_mat_add[9] * h2 +
                    ego_mat_add[10] * d2 + ego_mat_add[11];
        double w3 = ego_mat_add[0] * w2 + ego_mat_add[1] * h2 +
                    ego_mat_add[2] * d2 + ego_mat_add[3];
        double h3 = ego_mat_add[4] * w2 + ego_mat_add[5] * h2 +
                    ego_mat_add[6] * d2 + ego_mat_add[7];
        auto coors_x = (w3 - coors_x_min) / voxel_x;
        auto coors_y = (h3 - coors_y_min) / voxel_y;
        auto coors_z = (d3 - coors_z_min) / voxel_z;

        key_value[k * width * height + i * width + j].xf =
            static_cast<float>(coors_x);
        key_value[k * width * height + i * width + j].yf =
            static_cast<float>(coors_y);
        key_value[k * width * height + i * width + j].zf =
            static_cast<float>(coors_z);
      }
    }
  }
  return 0;
}

int32_t coor_reduce_hdim_cpu(CustomCoorPoint *o_coor_value,
                             const CustomCoorPoint *i_coor_value, int32_t bs,
                             int32_t nc, int32_t dz, int32_t height,
                             int32_t width) {
  int32_t count_y = bs * nc * dz;
  for (int32_t i{0}; i < count_y; ++i) {
    for (int32_t j{0}; j < width; ++j) {
      for (int32_t k{0}; k < height; ++k) {
        o_coor_value[(i * width + j)].xf +=
            i_coor_value[(i * height * width + k * width + j)].xf;
        o_coor_value[(i * width + j)].yf +=
            i_coor_value[(i * height * width + k * width + j)].yf;
        o_coor_value[(i * width + j)].zf +=
            i_coor_value[(i * height * width + k * width + j)].zf;
      }
      o_coor_value[(i * width + j)].x =
          int32_t(o_coor_value[(i * width + j)].xf / float(height));
      o_coor_value[(i * width + j)].y =
          int32_t(o_coor_value[(i * width + j)].yf / float(height));
      o_coor_value[(i * width + j)].z =
          int32_t(o_coor_value[(i * width + j)].zf / float(height));
    }
  }
  return 0;
}

int32_t coor_permute_cpu(CustomCoorPoint *o_coor_value,
                         const CustomCoorPoint *i_coor_value, int32_t *idx_ptr,
                         int32_t bs, int32_t nc, int32_t dz, int32_t width,
                         int32_t out_height, int32_t out_width,
                         int32_t invalid_x, int32_t invalid_y) {
  const int32_t nthreads = bs * nc * dz * width;
  for (int32_t index{0}; index < nthreads; ++index) {
    const int w = index % width;
    const int d = (index / width) % dz;
    const int c = (index / width / dz) % nc;
    const int b = index / width / dz / nc;
    const int new_index = b * nc * dz * width + c * dz * width + w * dz + d;
    //  || i_coor_value[index].z < 0 || i_coor_value[index].z >= 1
    if (i_coor_value[index].x < 0 || i_coor_value[index].x > out_width - 1 ||
        i_coor_value[index].y < 0 || i_coor_value[index].y > out_height - 1) {
      // invalid points
      o_coor_value[new_index].x = invalid_x;
      o_coor_value[new_index].y = invalid_y;
    } else {
      o_coor_value[new_index].x = i_coor_value[index].x;
      o_coor_value[new_index].y = i_coor_value[index].y;
      o_coor_value[new_index].z = i_coor_value[index].z;
    }
    idx_ptr[new_index] =
        o_coor_value[new_index].y * out_width + o_coor_value[new_index].x;
  }
  return 0;
}

bool gen_aug_ego_pose(std::vector<double> &augMat, std::vector<double> &egoMat,
                      std::vector<double> const &Cam_Intrins,
                      std::vector<double> const &Cam_Extrins,
                      std::vector<float> const &resize_scale,
                      std::vector<CamCrop> const &rect) {
  auto n_cam = static_cast<int32_t>(rect.size());
  // gen aug_pose
  for (int32_t i{0}; i < n_cam; ++i) {
    CamCrop CropArea = rect[i];
    float ResizeScale = resize_scale[i];
    Eigen::Matrix4d Pose1 = Eigen::Matrix4d::Identity();
    Pose1 << ResizeScale, 0, 0, -static_cast<float>(CropArea.x), 0, ResizeScale,
        0, -static_cast<float>(CropArea.y), 0, 0, 1, 0, 0, 0, 0, 1;
    Eigen::Matrix4d Pose1_Inv_T = Pose1.inverse().transpose();
    for (int32_t j{0}; j < 16; ++j) {
      augMat[i * 16 + j] = Pose1_Inv_T(j);
    }
  }

  // gen ego_pose
  for (int32_t i{0}; i < n_cam; ++i) {
    Eigen::Matrix3d Intrins;
    Intrins << Cam_Intrins[i * 9 + 0], Cam_Intrins[i * 9 + 1],
        Cam_Intrins[i * 9 + 2], Cam_Intrins[i * 9 + 3], Cam_Intrins[i * 9 + 4],
        Cam_Intrins[i * 9 + 5], Cam_Intrins[i * 9 + 6], Cam_Intrins[i * 9 + 7],
        Cam_Intrins[i * 9 + 8];

    Eigen::Matrix3d Extrins_Rot;
    Extrins_Rot << Cam_Extrins[i * 16 + 0], Cam_Extrins[i * 16 + 1],
        Cam_Extrins[i * 16 + 2], Cam_Extrins[i * 16 + 4],
        Cam_Extrins[i * 16 + 5], Cam_Extrins[i * 16 + 6],
        Cam_Extrins[i * 16 + 8], Cam_Extrins[i * 16 + 9],
        Cam_Extrins[i * 16 + 10];

    Eigen::Matrix3d Combine = (Extrins_Rot * Intrins.inverse()).transpose();
    egoMat[i * 16 + 0] = Combine(0);
    egoMat[i * 16 + 1] = Combine(1);
    egoMat[i * 16 + 2] = Combine(2);
    egoMat[i * 16 + 3] = Cam_Extrins[i * 16 + 3];
    egoMat[i * 16 + 4] = Combine(3);
    egoMat[i * 16 + 5] = Combine(4);
    egoMat[i * 16 + 6] = Combine(5);
    egoMat[i * 16 + 7] = Cam_Extrins[i * 16 + 7];
    egoMat[i * 16 + 8] = Combine(6);
    egoMat[i * 16 + 9] = Combine(7);
    egoMat[i * 16 + 10] = Combine(8);
    egoMat[i * 16 + 11] = Cam_Extrins[i * 16 + 11];
    egoMat[i * 16 + 12] = 0;
    egoMat[i * 16 + 13] = 0;
    egoMat[i * 16 + 14] = 0;
    egoMat[i * 16 + 15] = 1;
  }
  return true;
}
}  // namespace perception
}  // namespace autodrive
